package com.ftf.tp.blog.infrastructure.config;

import co.elastic.clients.elasticsearch._types.FieldValue;
import co.elastic.clients.elasticsearch._types.SortOptions;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.query_dsl.BoolQuery;
import co.elastic.clients.elasticsearch._types.query_dsl.QueryBuilders;
import co.elastic.clients.elasticsearch.core.search.SourceConfig;
import co.elastic.clients.elasticsearch.core.search.SourceConfigBuilders;
import lombok.Getter;


import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author tanghao
 * @date 2022/8/19
 */
@Getter
public class EsQueryWrapper {

    private final BoolQuery.Builder boolBuilder = new BoolQuery.Builder();

    private List<SortOptions> sortList = new ArrayList<>();

    private Integer size = 10;

    private Integer from = 0;

    private List<EsAggWrapper> aggWrapperList;

    private final SourceConfig.Builder sourceConfigBuilders = new SourceConfig.Builder();

    public EsQueryWrapper() {
    }

    public EsQueryWrapper(Integer size, Integer from, List<SortOptions> sortList) {
        this.size = size;
        this.from = from;
        this.sortList = sortList;
    }

    /**
     * 聚合标志位
     */
    protected boolean aggFlag = false;

    public BoolQuery.Builder getBoolBuilder() {
        return boolBuilder;
    }

    public SourceConfig.Builder getSourceConfigBuilders() {
        return sourceConfigBuilders;
    }

    public EsQueryWrapper eq(String field, String value) {
        boolBuilder.must(m -> m.term(t -> t.field(field).value(value))).build()._toQuery();
        return this;
    }

    public EsQueryWrapper eq(boolean condition, String field, String value) {
        if (condition) {
            boolBuilder.filter(m -> m.term(t -> t.field(field).value(value)));
            return this;
        }
        return eq(field, value);
    }

    public EsQueryWrapper ne(boolean condition, String field, String value) {
        if (condition) {
            boolBuilder.filter(m -> m.term(t -> t.field(field).value(value)));
            return this;
        }
        return eq(field, value);
    }

    public EsQueryWrapper in(String field, List<String> valueList) {
        List<FieldValue> fieldValues = valueList.stream()
                .map(value -> FieldValue.of(f -> f.stringValue(value))).collect(Collectors.toList());
        boolBuilder.must(m -> m.terms(t -> t.field(field).terms(ts -> ts.value(fieldValues))));
        return this;
    }

    public EsQueryWrapper in(boolean condition, String field, List<String> valueList) {
        if (condition) {
            List<FieldValue> fieldValues = valueList.stream()
                    .map(value -> FieldValue.of(f -> f.stringValue(value))).collect(Collectors.toList());
            boolBuilder.filter(m -> m.terms(t -> t.field(field).terms(ts -> ts.value(fieldValues))));
            return this;
        }
        return in(field, valueList);
    }

    public EsQueryWrapper match(String field, String value) {
        boolBuilder.must(m -> m.match(t -> t.field(field).query(value)));
        return this;
    }

    public EsQueryWrapper match(boolean condition, String field, String value) {
        if (condition) {
            boolBuilder.filter(m -> m.match(t -> t.field(field).query(value)));
        }
        return this;
    }

    /**
     * not嵌套 不会计算分数
     */
    public EsQueryWrapper not(Consumer<EsQueryWrapper> consumer) {
        EsQueryWrapper queryWrapper = new EsQueryWrapper();
        consumer.accept(queryWrapper);
        boolBuilder.mustNot(queryWrapper.boolBuilder.build()._toQuery());
        return this;
    }

    public EsQueryWrapper matchPhrase(String field, String value) {
        boolBuilder.must(m -> m.matchPhrase(t -> t.field(field).query(value)));
        return this;
    }

    public EsQueryWrapper matchPhrase(boolean condition, String field, String value) {
        if (condition) {
            boolBuilder.filter(m -> m.matchPhrase(t -> t.field(field).query(value)));
        }
        return this;
    }

    public EsQueryWrapper or(Consumer<EsQueryWrapper> consumer) {
        EsQueryWrapper queryWrapper = new EsQueryWrapper();
        consumer.accept(queryWrapper);
        boolBuilder.should(queryWrapper.boolBuilder.build()._toQuery());
        return this;
    }

    public EsQueryWrapper and(Consumer<EsQueryWrapper> consumer) {
        EsQueryWrapper queryWrapper = new EsQueryWrapper();
        consumer.accept(queryWrapper);
        boolBuilder.must(queryWrapper.boolBuilder.build()._toQuery());
        return this;
    }

    public EsQueryWrapper or() {
        EsQueryWrapper queryWrapper = new EsQueryWrapper(size, from, sortList);
        queryWrapper.boolBuilder.should(boolBuilder.build()._toQuery());
        return queryWrapper;
    }

    public EsQueryWrapper and() {
        EsQueryWrapper queryWrapper = new EsQueryWrapper(size, from, sortList);
        queryWrapper.boolBuilder.must(this.boolBuilder.build()._toQuery());
        return queryWrapper;
    }


    public EsQueryWrapper order(String field, SortOrder sortOrder) {
        sortList.add(SortOptions.of(s -> s.field(f -> f.field(field).order(sortOrder))));
        return this;
    }

    public EsQueryWrapper size(Integer size) {
        size = this.size;
        return this;
    }

    public EsQueryWrapper from(Integer from) {
        from = this.from;
        return this;
    }

    public EsQueryWrapper excludes(List<String> fieldList) {
        sourceConfigBuilders.filter(s -> s.excludes(fieldList));
        return this;
    }

    public EsQueryWrapper excludes(String value, String... values) {
        sourceConfigBuilders.filter(s -> s.excludes(value, values));
        return this;
    }

    public EsQueryWrapper includes(List<String> fieldList) {
        sourceConfigBuilders.filter(s -> s.includes(fieldList));
        return this;
    }

    public EsQueryWrapper includes(String value, String... values) {
        sourceConfigBuilders.filter(s -> s.includes(value, values));
        return this;
    }


    public EsAggWrapper aggTerms(String aggName, String field) {
        aggFlag = true;
        EsAggWrapper aggWrapper = new EsAggWrapper().aggTermsInit(aggName, field, 1);
        aggWrapperList.add(aggWrapper);
        return aggWrapper;
    }

    public EsAggWrapper aggTerms(String aggName, String field, Function<EsAggWrapper, EsAggWrapper> fn) {
        aggFlag = true;
        EsAggWrapper aggWrapper = new EsAggWrapper().aggTermsInit(aggName, field, 1);

        EsAggWrapper childWrapper = new EsAggWrapper();
        EsAggWrapper apply = fn.apply(childWrapper);
        aggWrapper.getEsAggWrapperList().add(apply);
        aggWrapperList.add(aggWrapper);
        return aggWrapper;
    }


}
